import os
import numpy as np
import cv2


def upsample_flow(flow, scale_factor, interpolation=cv2.INTER_CUBIC):
    h, w = flow.shape[:2]
    new_h, new_w = int(h * scale_factor), int(w * scale_factor)
    flow_up_x = cv2.resize(flow[:, :, 0], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up_y = cv2.resize(flow[:, :, 1], (new_w, new_h), interpolation=interpolation) * scale_factor
    flow_up = np.stack((flow_up_x, flow_up_y), axis=2)
    return flow_up


def process_and_save_upsampled_flow(input_dir, output_dir, scale_factor):
    os.makedirs(output_dir, exist_ok=True)

    for filename in os.listdir(input_dir):
        if filename.endswith('.npy'):
            input_flow_path = os.path.join(input_dir, filename)
            flow = np.load(input_flow_path)
            flow_upsampled = upsample_flow(flow, scale_factor)

            output_flow_path = os.path.join(output_dir, filename)
            np.save(output_flow_path, flow_upsampled)
            print(f"Upsampled and saved {filename} to {output_flow_path}")


# 示例用法
input_dir = 'E:\\wafer52\\32nm_flow_fine_align(3)'  # 替换为实际输入目录
output_dir = 'E:\\wafer52\\32nm_flow_fine_align_up_to_8nm(3)'  # 替换为实际输出目录

# 上采样因子
scale_factor = 4096 / 1024

process_and_save_upsampled_flow(input_dir, output_dir, scale_factor)

print(f"All flows upsampled and saved to {output_dir}")
